Write a custom CUDA kernel to optimize `ArcFace Loss` (Additive Angular Margin Loss).

Formula: Loss = -log( exp(s * cos(theta_yi + m)) / Sum(exp(s * cos(theta_j_modified))) )
Where:
- `cos(theta)` is the input cosine matrix `x`.
- For target class `yi`: we replace `cos(theta_yi)` with `cos(theta_yi + m)`.
- For other classes: use `cos(theta_j)` as is.

Problem Analysis:
1. Expensive Transcendental Functions: A standard implementation computes `acos(x)` to get theta, adds `m`, and then computes `cos(theta + m)`. `acos` and `cos` are computationally expensive operations on the GPU (Special Function Units).
2. Memory Bandwidth: Creating intermediate masks/indices to modify only the target class requires significant memory traffic.
3. Operator Chaining: The sequence `acos -> add -> cos -> scale -> softmax -> cross_entropy` creates multiple kernel launches and intermediate tensors.

Optimization Strategy: Fused Trigonometric Identity Kernel

1. Trigonometric Expansion: Instead of calculating angles, use the subtraction formula: `cos(theta + m) = cos(theta)cos(m) - sin(theta)sin(m)`.
   Since `cos(theta) = x` and `sin(theta) = sqrt(1 - x^2)`, the target logit becomes:
   `target_logit = s * (x * cos_m - sqrt(1 - x^2) * sin_m)`.
   This avoids expensive `acos` and `cos` instructions completely.

2. One-Block-per-Row: Launch one thread block per sample.

3. On-the-Fly Logic: The kernel iterates through the class dimension. It checks if the current index is the target.
   - If Target: Apply the optimized formula above.
   - If Not Target: Simply apply `s * x`.

4. Fused Reduction: Perform Online Softmax (find max, sum exp) within the same pass using Shared Memory reductions, then compute the final NLL loss.

Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

BATCH_SIZE = 512
NUM_CLASSES = 10000 
SHAPE = (BATCH_SIZE, NUM_CLASSES)

# ArcFace 超参数
SCALE_S = 64.0
MARGIN_M = 0.50

class ArcFaceLoss(nn.Module):
    """
    Standard PyTorch implementation of ArcFace.
    """
    def __init__(self, s=64.0, m=0.50, reduction='mean'):
        super(ArcFaceLoss, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction
        
        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        # 阈值，防止 acos 出界
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        # cosine: (N, C)
        # label: (N)
        
        # 1. 提取 target 对应的 cosine 值
        pick_cosine = cosine.gather(1, label.view(-1, 1)).squeeze(1)
        
        # 2. 计算 theta: acos(x)
        theta = torch.acos(torch.clamp(pick_cosine, -1.0 + 1e-7, 1.0 - 1e-7))
        
        # 3. 加 margin: cos(theta + m)
        target_logit = torch.cos(theta + self.m)
        
        # 4. 替换回原矩阵 (这里通常非常慢)
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, label.view(-1, 1), 1.0)
        
        # output = (1-one_hot) * cosine + one_hot * target_logit
        logits = cosine * (1.0 - one_hot) + target_logit.unsqueeze(1) * one_hot
        
        # 5. Scale
        logits = logits * self.s
        
        # 6. CrossEntropy
        loss = F.cross_entropy(logits, label, reduction='none')
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, s=64.0, m=0.50, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = ArcFaceLoss(s=s, m=m, reduction=reduction)
    
    def forward(self, cosine, label):
        return self.loss_fn(cosine, label)

def get_inputs():
    cosine = torch.randn(SHAPE, dtype=torch.float32)
    cosine = torch.clamp(cosine, -0.99, 0.99)
    label = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [cosine.contiguous(), label.contiguous()]

def get_init_inputs():
    return [SCALE_S, MARGIN_M, 'none']